#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A股股票智能筛选系统
实时联网获取数据并根据多维度指标筛选优质股票
"""

import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import List, Dict, Tuple
import warnings
warnings.filterwarnings('ignore')

# 配置参数
class ScreenerConfig:
    # 基础条件 - 已优化调整
    MIN_PRICE = 7
    MAX_PRICE = 50
    MIN_VOLUME_20D = 30_000_000  # 3000万元（降低）
    MIN_LIST_YEARS = 1.5  # 1.5年（放宽）
    MIN_MARKET_CAP = 20_00_000_000  # 20亿元（大幅降低）
    
    # 基本面指标
    MIN_PE = 5
    MAX_PE = 35  # 放宽
    MIN_PB = 0.8
    MAX_PB = 6  # 放宽
    MIN_ROE = 8  # 百分比（降低）
    MAX_DEBT_RATIO = 70  # 百分比（放宽）
    MIN_REVENUE_GROWTH = 10  # 百分比（降低）
    
    # 技术面指标
    MIN_3M_RETURN = -15  # 百分比（放宽）
    MAX_3M_RETURN = 25  # 百分比（放宽）
    MAX_52W_DRAWDOWN = 40  # 百分比（放宽）
    
    # 优先行业关键词
    PRIORITY_INDUSTRIES = [
        '半导体', '集成电路', '软件', '云计算', '人工智能', 'AI',
        '创新药', 'CXO', '医疗器械', '医疗服务',
        '锂电池', '光伏', '储能', '新能源车', '电池',
        '工业自动化', '精密仪器', '新材料',
        '品牌消费', '连锁', '食品饮料', '白酒'
    ]
    
    OUTPUT_COUNT = 20


class StockDataFetcher:
    """股票数据获取类"""
    
    def __init__(self):
        self.fetch_time = datetime.now()
        print(f"[{self.fetch_time.strftime('%Y-%m-%d %H:%M:%S')}] 开始获取数据...")
    
    def get_all_stocks(self) -> pd.DataFrame:
        """获取所有A股股票基本信息"""
        print("→ 获取股票基本信息...")
        try:
            # 获取沪深京A股实时行情
            df = ak.stock_zh_a_spot_em()
            print(f"  共获取 {len(df)} 支股票")
            return df
        except Exception as e:
            print(f"  ✗ 获取失败: {e}")
            return pd.DataFrame()
    
    def get_stock_individual_info(self, symbol: str) -> Dict:
        """获取个股详细信息"""
        try:
            info = ak.stock_individual_info_em(symbol=symbol)
            return dict(zip(info['item'], info['value']))
        except:
            return {}
    
    def get_financial_data(self, symbol: str) -> Dict:
        """获取财务指标（尝试多个数据源）"""
        try:
            # 方法1: 尝试获取财务分析指标
            df = ak.stock_financial_analysis_indicator(symbol=symbol)
            if not df.empty:
                latest = df.iloc[0]
                return latest.to_dict()
        except:
            pass
        
        try:
            # 方法2: 尝试获取实时指标
            df = ak.stock_a_lg_indicator(symbol=symbol)
            if not df.empty:
                latest = df.iloc[0]
                return latest.to_dict()
        except:
            pass
        
        return {}
    
    def get_historical_prices(self, symbol: str, days: int = 250) -> pd.DataFrame:
        """获取历史行情数据"""
        try:
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=days+30)).strftime('%Y%m%d')
            df = ak.stock_zh_a_hist(symbol=symbol, period="daily", 
                                   start_date=start_date, end_date=end_date, adjust="qfq")
            return df
        except:
            return pd.DataFrame()


class StockScreener:
    """股票筛选核心类"""
    
    def __init__(self):
        self.fetcher = StockDataFetcher()
        self.config = ScreenerConfig()
        self.stats = {
            '初始股票池': 0,
            '基础条件过滤后': 0,
            '基本面筛选后': 0,
            '行业筛选后': 0,
            '技术面筛选后': 0,
            '最终入选': 0
        }
    
    def run(self) -> pd.DataFrame:
        """执行完整筛选流程"""
        # 1. 获取所有股票
        all_stocks = self.fetcher.get_all_stocks()
        if all_stocks.empty:
            print("✗ 无法获取股票数据")
            return pd.DataFrame()
        
        self.stats['初始股票池'] = len(all_stocks)
        
        # 2. 基础条件过滤
        print("\n【阶段1】基础条件过滤...")
        stocks = self._filter_basic(all_stocks)
        self.stats['基础条件过滤后'] = len(stocks)
        print(f"  通过基础条件: {len(stocks)} 支")
        
        if stocks.empty:
            return pd.DataFrame()
        
        # 3. 获取详细数据并进行深度筛选
        print("\n【阶段2】深度数据获取与筛选...")
        results = []
        total = min(len(stocks), 500)  # 限制查询数量，避免API限流
        
        for idx, (_, stock) in enumerate(stocks.head(total).iterrows(), 1):
            if idx % 20 == 0:
                print(f"  处理进度: {idx}/{total}")
            
            stock_data = self._enrich_stock_data(stock)
            if stock_data:
                results.append(stock_data)
        
        if not results:
            print("✗ 未找到符合条件的股票")
            return pd.DataFrame()
        
        df_enriched = pd.DataFrame(results)
        
        # 4. 基本面筛选
        print("\n【阶段3】基本面筛选...")
        df_fundamental = self._filter_fundamental(df_enriched)
        self.stats['基本面筛选后'] = len(df_fundamental)
        print(f"  通过基本面筛选: {len(df_fundamental)} 支")
        
        # 5. 行业筛选
        print("\n【阶段4】行业筛选...")
        df_industry = self._filter_industry(df_fundamental)
        self.stats['行业筛选后'] = len(df_industry)
        print(f"  通过行业筛选: {len(df_industry)} 支")
        
        # 6. 技术面筛选
        print("\n【阶段5】技术面筛选...")
        df_technical = self._filter_technical(df_industry)
        self.stats['技术面筛选后'] = len(df_technical)
        print(f"  通过技术面筛选: {len(df_technical)} 支")
        
        # 7. 综合评分与排序
        df_final = self._score_and_rank(df_technical)
        df_final = df_final.head(self.config.OUTPUT_COUNT)
        self.stats['最终入选'] = len(df_final)
        
        return df_final
    
    def _filter_basic(self, df: pd.DataFrame) -> pd.DataFrame:
        """基础条件过滤"""
        # 排除ST股票和北交所
        df = df[~df['名称'].str.contains('ST|退', na=False)]
        df = df[~df['代码'].str.startswith('8')]  # 北交所代码以8开头
        df = df[~df['代码'].str.startswith('4')]  # 北交所代码以4开头
        
        # 股价区间
        df = df[(df['最新价'] >= self.config.MIN_PRICE) & 
                (df['最新价'] <= self.config.MAX_PRICE)]
        
        # 成交额过滤（使用成交额字段）
        if '成交额' in df.columns:
            df = df[df['成交额'] >= self.config.MIN_VOLUME_20D]
        
        # 总市值过滤
        if '总市值' in df.columns:
            df = df[df['总市值'] >= self.config.MIN_MARKET_CAP]
        
        return df.reset_index(drop=True)
    
    def _enrich_stock_data(self, stock: pd.Series) -> Dict:
        """丰富单个股票数据"""
        try:
            symbol = stock['代码']
            data = {
                '股票代码': symbol,
                '股票名称': stock['名称'],
                '当前价格': stock['最新价'],
                '涨跌幅': stock.get('涨跌幅', 0),
                '换手率': stock.get('换手率', 0),
                '总市值': stock.get('总市值', 0),
                '流通市值': stock.get('流通市值', 0),
                '市盈率': stock.get('市盈率-动态', np.nan),
                '市净率': stock.get('市净率', np.nan),
            }
            
            # 获取个股信息
            info = self.fetcher.get_stock_individual_info(symbol)
            if info:
                data['所属行业'] = info.get('行业', '未知')
                data['上市时间'] = info.get('上市时间', '')
                
                # 检查上市时间
                if data['上市时间']:
                    try:
                        list_date = datetime.strptime(data['上市时间'], '%Y-%m-%d')
                        years_since_list = (datetime.now() - list_date).days / 365
                        if years_since_list < self.config.MIN_LIST_YEARS:
                            return None
                    except:
                        pass
            
            # 获取财务数据
            financial = self.fetcher.get_financial_data(symbol)
            if financial:
                data['净资产收益率'] = financial.get('净资产收益率', np.nan)
                data['资产负债率'] = financial.get('资产负债率', np.nan)
                data['营业收入同比增长'] = financial.get('营业收入同比增长', np.nan)
                data['净利润'] = financial.get('净利润', np.nan)
            
            # 获取历史行情
            hist = self.fetcher.get_historical_prices(symbol)
            if not hist.empty and len(hist) >= 60:
                data['60日均价'] = hist['收盘'].tail(60).mean()
                data['52周最高'] = hist['收盘'].tail(250).max() if len(hist) >= 250 else hist['收盘'].max()
                
                # 计算3个月涨跌幅
                if len(hist) >= 60:
                    price_60d_ago = hist.iloc[-60]['收盘']
                    data['3个月涨跌幅'] = ((data['当前价格'] - price_60d_ago) / price_60d_ago) * 100
                
                # 计算52周回撤
                data['52周回撤'] = ((data['52周最高'] - data['当前价格']) / data['52周最高']) * 100
            
            return data
        except Exception as e:
            return None
    
    def _filter_fundamental(self, df: pd.DataFrame) -> pd.DataFrame:
        """基本面筛选"""
        df = df.dropna(subset=['市盈率', '市净率', '净资产收益率'])
        
        # PE筛选
        df = df[(df['市盈率'] >= self.config.MIN_PE) & 
                (df['市盈率'] <= self.config.MAX_PE)]
        
        # PB筛选
        df = df[(df['市净率'] >= self.config.MIN_PB) & 
                (df['市净率'] <= self.config.MAX_PB)]
        
        # ROE筛选
        df = df[df['净资产收益率'] >= self.config.MIN_ROE]
        
        # 资产负债率筛选
        df = df[df['资产负债率'] <= self.config.MAX_DEBT_RATIO]
        
        # 营收增长筛选
        df = df[df['营业收入同比增长'] >= self.config.MIN_REVENUE_GROWTH]
        
        return df.reset_index(drop=True)
    
    def _filter_industry(self, df: pd.DataFrame) -> pd.DataFrame:
        """行业筛选"""
        def is_priority_industry(industry: str) -> bool:
            if pd.isna(industry):
                return False
            return any(keyword in industry for keyword in self.config.PRIORITY_INDUSTRIES)
        
        df = df[df['所属行业'].apply(is_priority_industry)]
        return df.reset_index(drop=True)
    
    def _filter_technical(self, df: pd.DataFrame) -> pd.DataFrame:
        """技术面筛选"""
        # 3个月涨跌幅
        df = df[(df['3个月涨跌幅'] >= self.config.MIN_3M_RETURN) & 
                (df['3个月涨跌幅'] <= self.config.MAX_3M_RETURN)]
        
        # 52周回撤
        df = df[df['52周回撤'] <= self.config.MAX_52W_DRAWDOWN]
        
        # MA60支撑
        df = df[df['当前价格'] >= df['60日均价']]
        
        return df.reset_index(drop=True)
    
    def _score_and_rank(self, df: pd.DataFrame) -> pd.DataFrame:
        """综合评分与排序"""
        # 简单评分模型：ROE权重40%，营收增长30%，估值合理性30%
        df['ROE得分'] = df['净资产收益率'] / df['净资产收益率'].max() * 40
        df['增长得分'] = df['营业收入同比增长'] / df['营业收入同比增长'].max() * 30
        df['估值得分'] = (1 - (df['市盈率'] - self.config.MIN_PE) / 
                         (self.config.MAX_PE - self.config.MIN_PE)) * 30
        
        df['综合得分'] = df['ROE得分'] + df['增长得分'] + df['估值得分']
        
        # 生成推荐理由
        df['推荐理由'] = df.apply(self._generate_reason, axis=1)
        
        # 排序
        df = df.sort_values('综合得分', ascending=False)
        
        return df
    
    def _generate_reason(self, row: pd.Series) -> str:
        """生成推荐理由"""
        reasons = []
        
        if row['净资产收益率'] >= 20:
            reasons.append('ROE优秀')
        elif row['净资产收益率'] >= 15:
            reasons.append('ROE良好')
        
        if row['营业收入同比增长'] >= 30:
            reasons.append('高速增长')
        elif row['营业收入同比增长'] >= 20:
            reasons.append('稳定增长')
        
        if row['市盈率'] < 15:
            reasons.append('估值偏低')
        elif row['市盈率'] < 20:
            reasons.append('估值合理')
        
        if row['资产负债率'] < 30:
            reasons.append('低负债')
        
        industry_str = f"{row['所属行业']}行业"
        
        return f"{industry_str}，" + '，'.join(reasons)


class ReportGenerator:
    """报告生成类"""
    
    def __init__(self, df: pd.DataFrame, stats: Dict, fetch_time: datetime):
        self.df = df
        self.stats = stats
        self.fetch_time = fetch_time
    
    def print_console(self):
        """控制台输出"""
        from tabulate import tabulate
        
        print("\n" + "="*100)
        print("A股股票智能筛选报告".center(100))
        print("="*100)
        
        # 数据时间戳
        print(f"\n【数据时间】")
        print(f"  数据获取时间: {self.fetch_time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"  最新交易日: {datetime.now().strftime('%Y-%m-%d')}")
        
        # 筛选统计
        print(f"\n【筛选统计】")
        for key, value in self.stats.items():
            print(f"  {key}: {value} 支")
        
        if self.stats['最终入选'] > 0:
            rate = (self.stats['最终入选'] / self.stats['初始股票池']) * 100
            print(f"  筛选率: {rate:.2f}%")
        
        # 结果表格
        if not self.df.empty:
            print(f"\n【筛选结果】Top {len(self.df)} 支股票\n")
            
            display_df = self.df[[
                '股票代码', '股票名称', '所属行业', '当前价格', 
                '市盈率', '市净率', '净资产收益率', '资产负债率',
                '营业收入同比增长', '3个月涨跌幅', '52周回撤', '推荐理由'
            ]].copy()
            
            # 格式化数值
            display_df['当前价格'] = display_df['当前价格'].apply(lambda x: f"{x:.2f}")
            display_df['市盈率'] = display_df['市盈率'].apply(lambda x: f"{x:.2f}")
            display_df['市净率'] = display_df['市净率'].apply(lambda x: f"{x:.2f}")
            display_df['净资产收益率'] = display_df['净资产收益率'].apply(lambda x: f"{x:.1f}%")
            display_df['资产负债率'] = display_df['资产负债率'].apply(lambda x: f"{x:.1f}%")
            display_df['营业收入同比增长'] = display_df['营业收入同比增长'].apply(lambda x: f"{x:.1f}%")
            display_df['3个月涨跌幅'] = display_df['3个月涨跌幅'].apply(lambda x: f"{x:.1f}%")
            display_df['52周回撤'] = display_df['52周回撤'].apply(lambda x: f"{x:.1f}%")
            
            print(tabulate(display_df, headers='keys', tablefmt='grid', showindex=False))
        
        # 风险提示
        print(f"\n【风险提示】")
        print("  1. 本筛选结果仅供参考，不构成投资建议")
        print("  2. 股市有风险，投资需谨慎")
        print("  3. 历史业绩不代表未来表现")
        print("  4. 请结合自身风险承受能力谨慎决策")
        print("="*100 + "\n")
    
    def export_excel(self, filename: str = None):
        """导出Excel文件"""
        if self.df.empty:
            print("✗ 无数据可导出")
            return
        
        if filename is None:
            filename = f"A股筛选结果_{self.fetch_time.strftime('%Y%m%d_%H%M%S')}.xlsx"
        
        try:
            with pd.ExcelWriter(filename, engine='openpyxl') as writer:
                # Sheet1: 筛选结果
                output_df = self.df[[
                    '股票代码', '股票名称', '所属行业', '当前价格', '总市值',
                    '市盈率', '市净率', '净资产收益率', '资产负债率',
                    '营业收入同比增长', '3个月涨跌幅', '52周回撤', 
                    '综合得分', '推荐理由'
                ]].copy()
                
                output_df.to_excel(writer, sheet_name='筛选结果', index=False)
                
                # Sheet2: 详细数据
                self.df.to_excel(writer, sheet_name='详细数据', index=False)
                
                # Sheet3: 参数说明
                params_data = {
                    '参数名称': [
                        '股价区间', '日均成交额', '总市值', '上市时间',
                        'PE区间', 'PB区间', 'ROE', '资产负债率', '营收增长',
                        '3个月涨跌幅', '52周回撤', '优先行业'
                    ],
                    '标准值': [
                        f"{ScreenerConfig.MIN_PRICE}-{ScreenerConfig.MAX_PRICE}元",
                        f">{ScreenerConfig.MIN_VOLUME_20D/10000:.0f}万元",
                        f">{ScreenerConfig.MIN_MARKET_CAP/100000000:.0f}亿元",
                        f">{ScreenerConfig.MIN_LIST_YEARS}年",
                        f"{ScreenerConfig.MIN_PE}-{ScreenerConfig.MAX_PE}",
                        f"{ScreenerConfig.MIN_PB}-{ScreenerConfig.MAX_PB}",
                        f">{ScreenerConfig.MIN_ROE}%",
                        f"<{ScreenerConfig.MAX_DEBT_RATIO}%",
                        f">{ScreenerConfig.MIN_REVENUE_GROWTH}%",
                        f"{ScreenerConfig.MIN_3M_RETURN}%~{ScreenerConfig.MAX_3M_RETURN}%",
                        f"<{ScreenerConfig.MAX_52W_DRAWDOWN}%",
                        '科技/医疗/新能源/高端制造/消费升级'
                    ]
                }
                pd.DataFrame(params_data).to_excel(writer, sheet_name='筛选参数', index=False)
                
                # Sheet4: 统计信息
                stats_data = {
                    '统计项': list(self.stats.keys()),
                    '数量': list(self.stats.values())
                }
                stats_df = pd.DataFrame(stats_data)
                stats_df.to_excel(writer, sheet_name='筛选统计', index=False)
            
            print(f"✓ Excel文件已导出: {filename}")
        except Exception as e:
            print(f"✗ 导出失败: {e}")


def main():
    """主函数"""
    print("="*80)
    print("A股股票智能筛选系统".center(80))
    print("="*80)
    
    # 执行筛选
    screener = StockScreener()
    result_df = screener.run()
    
    if result_df.empty:
        print("\n✗ 未找到符合条件的股票，请调整筛选参数后重试")
        return
    
    # 生成报告
    reporter = ReportGenerator(result_df, screener.stats, screener.fetcher.fetch_time)
    reporter.print_console()
    reporter.export_excel()
    
    print("✓ 筛选完成！")


if __name__ == "__main__":
    main()
